#!/usr/bin/env python
# 当使用roslaunch来启动python节点时，无法找到对应路径，所以借助rospkg来查找功能包的路径
import sys
import rospkg
# 初始化rospkg
rospack = rospkg.RosPack()
# 获取当前功能包的路径
current_rospkg_directory = rospack.get_path('boat_vision')
sys.path.append(f"{current_rospkg_directory}/detect_yolox")

import rospy

import cv2
import numpy as np
from PIL import Image

from yolo import YOLO

from sensor_msgs.msg import Image as sensorImg
from sensor_msgs.msg import CompressedImage as sensorCompressedImg
from cv_bridge import CvBridge, CvBridgeError
from geometry_msgs.msg import Point

from collections import deque

class Vision():
    def __init__(self, crop_flag=True):
        self.yolo_model = YOLO()
        self.crop_flag = crop_flag
        self.init_publisher()
        rospy.init_node("yolox_detect")
        self.camera_topic = rospy.get_param("~camera_topic", "/camera/color/image_raw")
        print(f"camera_topic is {self.camera_topic}")
        if self.camera_topic == "/camera/color/image_raw":
            self.sub_image = rospy.Subscriber(self.camera_topic, sensorImg, self.image_callback, queue_size=10)
        elif self.camera_topic == "/camera/color/image_raw/compressed":
            self.sub_image = rospy.Subscriber(self.camera_topic, sensorCompressedImg, self.image_callback, queue_size=10)
        else:
            rospy.logerr("Invalid camera topic specified.")
        # 设置队列主要是为了想让物体位置波动时，输出的坐标位置波动不要太快
        self.queue = deque(maxlen=10)

    def image_callback(self, msg):
        try:
            bridge = CvBridge()
            frame = self.process_image_msg(msg, bridge)
            detect_image, detect_crop_image, detect_image_box = self.yolo_model.detect_image(frame, crop=self.crop_flag)
            detect_image = np.array(detect_image)
            detect_crop_image = np.array(detect_crop_image)
            detect_image = bridge.cv2_to_imgmsg(detect_image, encoding="rgb8")
            detect_crop_image = bridge.cv2_to_imgmsg(detect_crop_image, encoding="rgb8")
            center_x = (detect_image_box[1] + detect_image_box[3]) / 2
            center_y = (detect_image_box[0] + detect_image_box[2]) / 2
            self.queue.append([center_x, center_y])
            # 计算center_x和center_y的总和
            sum_x = sum(element[0] for element in self.queue)
            sum_y = sum(element[1] for element in self.queue)
            # 计算平均值
            center_average_x = sum_x / len(self.queue)
            center_average_y = sum_y / len(self.queue)
            print(f"center_average_x = {center_average_x}, center_average_y = {center_average_y}")
            center_point = Point(x=center_average_x, y=center_average_y, z=1)
            self.pub_detect_image.publish(detect_image)
            self.pub_detect_crop_image.publish(detect_crop_image)
            self.pub_bounding_box_center.publish(center_point)
        except CvBridgeError as e:
            rospy.logerr("CV Bridge Exception: %s", e)

    def process_image_msg(self, msg, bridge):
        if isinstance(msg, sensorImg):
            frame = bridge.imgmsg_to_cv2(msg, desired_encoding='rgb8')
        elif isinstance(msg, sensorCompressedImg):
            np_arr = np.fromstring(msg.data, np.uint8)
            frame = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        else:
            rospy.logerr("Invalid image message type")
            return None
        frame = Image.fromarray(np.uint8(frame))
        return frame

    def init_publisher(self):
        self.pub_detect_image = rospy.Publisher("/camera/color/image_raw/yolox_image", sensorImg, queue_size=10)
        self.pub_detect_crop_image = rospy.Publisher("/camera/color/image_raw/yolox_crop_image", sensorImg, queue_size=10)
        self.pub_bounding_box_center = rospy.Publisher("/camera/color/image_raw/yolox_bounding_box_center", Point, queue_size=10)

if __name__ == "__main__":
    try:
        vision = Vision(crop_flag=True)  # 设置为True或False，根据需求调整
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
